#[macro_use]
extern crate serde_derive;

mod handler;
pub use handler::Handler;

pub mod server;

pub mod utils;

mod shutdown;

use tokio::io::AsyncReadExt;
use tokio::io::{self, AsyncRead, AsyncWrite, AsyncWriteExt};
use tracing::error;

/// Error returned by most functions.
///
/// When writing a real application, one might want to consider a specialized
/// error handling crate or defining an error type as an `enum` of causes.
/// However, for our example, using a boxed `std::error::Error` is sufficient.
///
/// For performance reasons, boxing is avoided in any hot path. For example, in
/// `parse`, a custom error `enum` is defined. This is because the error is hit
/// and handled during normal execution when a partial frame is received on a
/// socket. `std::error::Error` is implemented for `parse::Error` which allows
/// it to be converted to `Box<dyn std::error::Error>`
pub type Error = Box<dyn std::error::Error + Send + Sync>;

/// A specialized `Result` type for mini-redis operations.
///
/// This is defined as a convenience.
pub type Result<T> = std::result::Result<T, Error>;

pub mod consts {
    pub const SOCKS5_VERSION: u8 = 0x05;
    pub const SOCKS5_RESERVED: u8 = 0x00;

    pub const SOCKS5_ADDR_TYPE_IPV4: u8 = 0x01;
    pub const SOCKS5_ADDR_TYPE_DOMAIN_NAME: u8 = 0x03;
    pub const SOCKS5_ADDR_TYPE_IPV6: u8 = 0x04;
}

#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct User {
    pub username: String,
    password: String,
}

/// Client Authentication Methods
pub enum AuthMethods {
    /// No Authentication
    NoAuth = 0x00,
    /// Authenticate with gssapi
    GssApi = 0x01,
    /// Authenticate with a username / password
    UserPass = 0x02,
    /// Cannot authenticate
    NoMethods = 0xFF,
}

#[derive(Debug, PartialEq)]
pub enum ReplyCode {
    Successed = 0x00,
    Failure = 0x01,
    ConnectNotAllow = 0x02,
    NetworkUnreachable = 0x03,
    HostUnreachable = 0x04,
    ConnectionRefused = 0x05,
    TtlExpired = 0x06,
    CommandNotSupported = 0x07,
    AddressTypeNotSupported = 0x08,
}

#[derive(PartialEq)]
enum Socks5Command {
    TCPConnect = 0x01,
    TCPBind = 0x02,
    UDPAssociate = 0x03,
    UNKNOWN,
}

impl Socks5Command {
    /// Parse Byte to Command
    fn from(n: usize) -> Socks5Command {
        match n {
            1 => Socks5Command::TCPConnect,
            2 => Socks5Command::TCPBind,
            3 => Socks5Command::UDPAssociate,
            _ => Socks5Command::UNKNOWN,
        }
    }
}

#[derive(PartialEq)]
enum AddrType {
    /// IP V4 address: X'01'
    V4 = 0x01,
    /// DOMAINNAME: X'03'
    Domain = 0x03,
    /// IP V6 address: X'04'
    V6 = 0x04,
    UNKNOWN,
}

impl AddrType {
    /// Parse Byte to Command
    fn from(n: usize) -> AddrType {
        match n {
            1 => AddrType::V4,
            3 => AddrType::Domain,
            4 => AddrType::V6,
            _ => AddrType::UNKNOWN,
        }
    }
}

#[allow(dead_code)]
struct Socks5Req {
    pub version: u8,
    pub command: Socks5Command,
    pub addr_type: AddrType,
    pub addr: Vec<u8>,
    pub port: u16,
}

impl Socks5Req {
    /// Parse a SOCKS Req from a TcpStream
    async fn from_stream<T>(stream: &mut T) -> Result<Self>
    where
        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
    {
        // From rfc 1928 (S4), the SOCKS request is formed as follows:
        //
        //    +----+-----+-------+------+----------+----------+
        //    |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
        //    +----+-----+-------+------+----------+----------+
        //    | 1  |  1  | X'00' |  1   | Variable |    2     |
        //    +----+-----+-------+------+----------+----------+
        //
        // Where:
        //
        //      o  VER    protocol version: X'05'
        //      o  CMD
        //         o  CONNECT X'01'
        //         o  BIND X'02'
        //         o  UDP ASSOCIATE X'03'
        //      o  RSV    RESERVED
        //      o  ATYP   address type of following address
        //         o  IP V4 address: X'01'
        //         o  DOMAINNAME: X'03'
        //         o  IP V6 address: X'04'
        //      o  DST.ADDR       desired destination address
        //      o  DST.PORT desired destination port in network octet
        //         order

        let mut packet = [0u8; 4];
        stream.read_exact(&mut packet).await?;

        if packet[0] != consts::SOCKS5_VERSION {
            error!("Socks5Req::from_stream unsupported version: {}", packet[0]);
            stream.shutdown().await?;
        }

        // Get command
        let command = Socks5Command::from(packet[1] as usize);
        if command == Socks5Command::UNKNOWN {
            error!("Socks5Req::from_stream addr type not supported");
            stream.shutdown().await?;
            return Err("addr type not supported".into());
        }

        // DST.address
        let addr_type = AddrType::from(packet[3] as usize);
        if addr_type == AddrType::UNKNOWN {
            error!("Socks5Req::from_stream addr type not supported");
            stream.shutdown().await?;
            return Err("addr type not supported".into());
        }

        // Get Addr from addr_type and stream
        let addr: Vec<u8> = match addr_type {
            AddrType::Domain => {
                let mut dlen = [0u8; 1];
                stream.read_exact(&mut dlen).await?;
                let mut domain = vec![0u8; dlen[0] as usize];
                stream.read_exact(&mut domain).await?;
                domain
            }
            AddrType::V4 => {
                let mut addr = [0u8; 4];
                stream.read_exact(&mut addr).await?;
                addr.to_vec()
            }
            AddrType::V6 => {
                let mut addr = [0u8; 16];
                stream.read_exact(&mut addr).await?;
                addr.to_vec()
            }
            _ => { Vec::new() }
        };

        // read DST.port
        let mut port = [0u8; 2];
        stream.read_exact(&mut port).await?;

        // Merge two u8s into u16
        let port = (u16::from(port[0]) << 8) | u16::from(port[1]);

        Ok(Socks5Req {
            version: packet[0],
            command,
            addr_type,
            addr,
            port,
        })
    }
}

pub struct SocksReply {
    // From rfc 1928 (S6),
    // the server evaluates the request, and returns a reply formed as follows:
    //
    //    +----+-----+-------+------+----------+----------+
    //    |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
    //    +----+-----+-------+------+----------+----------+
    //    | 1  |  1  | X'00' |  1   | Variable |    2     |
    //    +----+-----+-------+------+----------+----------+
    //
    // Where:
    //
    //      o  VER    protocol version: X'05'
    //      o  REP    Reply field:
    //         o  X'00' succeeded
    //         o  X'01' general SOCKS server failure
    //         o  X'02' connection not allowed by ruleset
    //         o  X'03' Network unreachable
    //         o  X'04' Host unreachable
    //         o  X'05' Connection refused
    //         o  X'06' TTL expired
    //         o  X'07' Command not supported
    //         o  X'08' Address type not supported
    //         o  X'09' to X'FF' unassigned
    //      o  RSV    RESERVED
    //      o  ATYP   address type of following address
    //         o  IP V4 address: X'01'
    //         o  DOMAINNAME: X'03'
    //         o  IP V6 address: X'04'
    //      o  BND.ADDR       server bound address
    //      o  BND.PORT       server bound port in network octet order
    //
    buf: [u8; 10],
}

impl SocksReply {
    pub fn new(status: ReplyCode) -> Self {
        let buf = [
            // VER
            consts::SOCKS5_VERSION,
            // REP
            status as u8,
            // RSV
            consts::SOCKS5_RESERVED,
            // ATYP
            1,
            // BND.ADDR
            0,
            0,
            0,
            0,
            // BND.PORT
            0,
            0,
        ];
        Self { buf }
    }

    pub async fn send<T>(&self, stream: &mut T) -> io::Result<()>
    where
        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
    {
        stream.write_all(&self.buf[..]).await?;
        Ok(())
    }
}
